
# -----------------------------------------------------------
rm(list=ls())
library(plm)
library(gmm)
library(spatialreg)
library(spdep)
library(splm)

# -----------------------------------------------------------
# Define a function to calculate the spillover productivity
swfn <- function(w) {
  sw <- c()
  for(j in 1:N) {
    sw <- c(sw, sum(w[-j])/(N-1))
  }
  return(sw)
}
# -----------------------------------------------------------
# The following function re-arrange a matrix into a vector
long <- function(X) {
  X2 <- c()
  for(i in 1:nrow(X)) {
    X2 <- c(X2, X[i,])
  }
  return(X2)
}
# -----------------------------------------------------------

# ========================================================
# DGP

set.seed(123)

# Parameters
N <- 100; T <- 10; R <- 1000
d0 <- 0.2; d1 <- 0.55; d2 <- 0; d3 <- 0    # productivity parameters
ak <- 0.25; am <- 0.65    # production function parameters
b1 <- 0.8; b2 <- 0.1    # I policy function parameters
r0 <- 0.01; r1 <- 0.6; r2 <- 0.3     # G parameters
de <- c(0.05, 0.075, 0.1, 0.125, 0.15)   # depreciation rate

# Empty objects
INITIAL <- DFDK <- DFDM <- c()   # matrixes for the estimated elasticities
O1a <- O1ap <- O1ase <- O2a <- O2ap <- O2ase <- c()

# spatial weight matrix
NB <- matrix(1/(N-1), N, N); diag(NB) <- 0; 

# Start the loop
for(ii in 1:R) {
  
  # initial values
  w0 <- runif(N, 1, 3)
  K0 <- runif(N, 10, 200)
  G0 <- runif(N, 0, 1); 
  I0 <- K0^b1*exp(b2*w0)
  M0 <- (am*K0^ak*exp(w0))^(1/(1-am))

  # errors
  eta <- matrix(rnorm(N*T, 0, sqrt(0.07)), N, T)
  xi <- matrix(rnorm(N*T, 0, sqrt(0.04)), N, T)   
  eps <- matrix(rnorm(N*T, 0, sqrt(0.01)), N, T)
  
  # Evolution processes
  # We will first generate state variables, and then control variables whose values
  # depend on state variables.
  Y <- K <- M <- I <- W <- G <- matrix(NA, N, T)
  
  for(i in 1:T) {
    if(i==1) {
      W[,i] <- d0 + d1*w0 + d2*swfn(w0) +d3*G0 + xi[,i]
      K[,i] <- (1-de)*K0 + I0
      G[,i] <- r0 + r1*G0 + r2*W[,i] + eps[,i]; 
      I[,i] <- K[,i]^b1*exp(b2*W[,i])
      M[,i] <- (am*K[,i]^ak*exp(W[,i]))^(1/(1-am))
      Y[,i] <- K[,i]^ak*M[,i]^am*exp(W[,i]+eta[,i])
    }
    else {
      W[,i] <- d0 + d1*W[,(i-1)] + d2*swfn(W[,(i-1)]) +d3*G[,(i-1)] + xi[,i]
      K[,i] <- (1-de)*K[,(i-1)] + I[,(i-1)]
      G[,i] <- r0 + r1*G[,(i-1)] + r2*W[,i] + eps[,i]; 
      I[,i] <- K[,i]^b1*exp(b2*W[,i])
      M[,i] <- (am*K[,i]^ak*exp(W[,i]))^(1/(1-am))
      Y[,i] <- K[,i]^ak*M[,i]^am*exp(W[,i]+eta[,i])
    }
  }
  
  pm <- mean(exp(eta))
  data <- data.frame(id=rep(1:N, each = T), year=rep(1:T, N), Y=long(Y), K=long(K), M=long(M), G=long(G), sm=pm*long(M)/long(Y) )
  
  # ----------------------------------------------
  # Start the estimation
  # Step 1
  h.lnbe <- mean(log(data$sm)); h.eta <- - log(data$sm) + h.lnbe; h.e <- mean(exp(h.eta))
  h.betam <- exp(h.lnbe)/h.e
  data$h.eta <- h.eta
  
  # ----------------------------------------------
  data <- pdata.frame(data)
  data$K1 <- lag(data$K); data$M1 <- lag(data$M); data$G1 <- lag(data$G); data$G2 <- lag(data$G, 2)

  data$y <- log(data$Y); data$k <- log(data$K); data$m <- log(data$M)
  data$k1 <- log(data$K1); data$m1 <- log(data$M1);
  data <- data[!is.na(data$k1),]
  data <- as.data.frame(data)
  print(ii)
  
  # ----------------------------------------------
  # Step 2
  # Define function "phi" 
  phi <- function(beta) {
    (1-h.betam)*data$m1 - h.lnbe - log(1/pm) - beta[1]*data$k1 
  }
  
  # Optimization
  RHO.temp <- OBJ <- BK <- c()
  for(ik in seq(0.01, 1, by=0.01)) {
    y.star <- data$y - h.betam*data$m - ik*data$k 
    p1 <- as.vector(phi(ik))
    var.poly <- poly(cbind(p1), degree = 2, raw = TRUE)
    mod1 <- lm(y.star~var.poly)
    RHO.temp <- cbind(RHO.temp, mod1$coeff)
    OBJ <- cbind(OBJ, sum(mod1$residuals^2))
    BK <- cbind(BK, ik)
  }
  target <- which.min(OBJ)
  
  # results - elas
  h.betak <- BK[target]
  DFDK <- cbind(DFDK, h.betak); DFDM <- cbind(DFDM, h.betam)
  
  # Post-reg
  # Generate variables
  data$w <- data$y - h.betak*data$k - h.betam*data$m - data$h.eta
  w1 <- as.vector(phi(c(h.betak)))
  
  for(i in 1:nrow(data)) {
    s1.t <- (as.numeric(data$year)==as.numeric(data$year[i])); s1.t[i] <- 0
    s1.b <- sum(s1.t)
    s1 <- if(s1.b>0) s1.t/s1.b else rep(0, length(s1.t));
    data$sG2[i] <- sum(s1*data$G2, na.rm = TRUE)
    data$sw1[i] <- sum(s1*w1)
  }
  
  pval <- function(x) {summary(x)[["coefficients"]][,"Pr(>|t|)"]} 
  seval <- function(x) {summary(x)[["coefficients"]][,"Std. Error"]}
  
  ols <- lm(w~G1+sG2, data = data); O1a <- cbind(O1a, ols$coefficients[c("G1","sG2")])
  O1ap <- cbind(O1ap, pval(ols)[c("G1","sG2")]); O1ase <- cbind(O1ase, seval(ols)[c("G1","sG2")])
  # ols <- lm(w~w1+G1+sG2, data = data); O1b <- cbind(O1b, ols$coefficients[c("w1","G1","sG2")])
  # O1bp <- cbind(O1bp, pval(ols)[c("w1","G1","sG2")]); O1bse <- cbind(O1bse, seval(ols)[c("w1","G1","sG2")])
  
  ols <- lm(w~G1+sw1, data = data); O2a <- cbind(O2a, ols$coefficients[c("G1","sw1")])
  O2ap <- cbind(O2ap, pval(ols)[c("G1","sw1")]); O2ase <- cbind(O2ase, seval(ols)[c("G1","sw1")])
  # ols <- lm(w~w1+G1+sw1, data = data); O2b <- cbind(O2b, ols$coefficients[c("w1","G1","sw1")])
  # O2bp <- cbind(O2bp, pval(ols)[c("w1","G1","sw1")]); O2bse <- cbind(O2bse, seval(ols)[c("w1","G1","sw1")])
  
}

# Define functions
rmse <- function(x, truevalue=0.25) sqrt(mean((x-truevalue)^2))
mad <- function(x, truevalue=0.25) mean(abs(x-truevalue))
median2 <- function(x) {format(round(median(x), 5), nsmall = 5)}
mean2 <- function(x) {format(round(mean(x), 5), nsmall = 5)}

# k coefficient
n11 <- mean(DFDK); n12 <- median(DFDK); n13 <- sd(DFDK); n14 <- rmse(DFDK, 0.25); n15 <- mad(DFDK, 0.25)
Table <- rbind(c("mean","median","sd","rmse","mad"),c(n11, n12, n13, n14, n15))
print(Table, quote=FALSE, na.print = "")

# regression stat - mean 
Table <- matrix(NA, 4, 4); Table <- cbind(c(" ","G1","sw1","sG2"), Table)
Table[2,2] <- mean2(O1a[1,]); Table[4,2] <- mean2(O1a[2,])
Table[2,4] <- mean2(O2a[1,]); Table[3,4] <- mean2(O2a[2,])
print(Table, quote=FALSE, na.print = "")

# regression stat - rmse 
Table <- matrix(NA, 4, 4); Table <- cbind(c(" ","G1","sw1","sG2"), Table)
Table[2,2] <- rmse(O1a[1,],0); Table[4,2] <- rmse(O1a[2,],0)
Table[2,4] <- rmse(O2a[1,],0); Table[3,4] <- rmse(O2a[2,],0)
print(Table, quote=FALSE, na.print = "")

